package org.example;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;

public class LeaderElection implements Watcher {

    private ZooKeeper zookeeper;
    private String path = "/my_election";
    private String serverName;
    private boolean isLeader = false;

    public LeaderElection(String connectString, String serverName, final CountDownLatch countDownLatch) throws IOException {
        this.serverName = serverName;
        this.zookeeper = new ZooKeeper(connectString, 3000, new Watcher() {
            @Override
            public void process(WatchedEvent event) {
                if (event.getState() == Event.KeeperState.SyncConnected) {
                    System.out.println("连接成功");
                    countDownLatch.countDown();
                }
            }
        });

    }

    public void waitConnection(CountDownLatch countDownLatch) throws Exception {
        if (ZooKeeper.States.CONNECTING == zookeeper.getState()) {
            System.out.println("等待连接...");
            countDownLatch.await();
        }
    }

    public void startElection() throws Exception {
        System.out.println("当前状态：" + zookeeper.getState());
        this.zookeeper.register(this);
        // 创建选举路径
        try {
            // 尝试创建节点
            zookeeper.create(path, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        } catch (KeeperException.NodeExistsException e) {
            // 如果节点已存在，则获取节点状态或进行其他处理
            Stat stat = zookeeper.exists(path, false);
            System.out.println("Node already exists with stat: " + stat);
        } catch (InterruptedException | KeeperException e) {
            // 处理其他异常...
        }
        // 创建临时有序节点
        String nodePath = zookeeper.create(path + "/server-", serverName.getBytes(StandardCharsets.UTF_8),
                ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);

        // 获取所有子节点
        List<String> children = zookeeper.getChildren(path, true);

        Collections.sort(children);

        // 检查是否是最小的节点
        if (nodePath.substring(nodePath.lastIndexOf("/") + 1).equals(children.get(0))) {
            System.out.println(serverName + " 成为领导者！");
            zookeeper.exists("/abc", true);
            isLeader = true;
        } else {
            System.out.println(serverName + " 不是领导者。");
            String prevNode = children.get(Collections.binarySearch(children, nodePath.substring(nodePath.lastIndexOf("/") + 1)) - 1);
            watchNode(path + "/" + prevNode);
            try {
                zookeeper.setData("/abc", ("i am " + serverName).getBytes(StandardCharsets.UTF_8), -1);
            } catch (KeeperException.NodeExistsException e) {
                // 如果节点已存在，则获取节点状态或进行其他处理
                Stat stat = zookeeper.exists(path, false);
                System.out.println("Node already exists with stat: " + stat);
            } catch (InterruptedException | KeeperException e) {
                // 处理其他异常...
            }
        }
    }

    public void watchNode(final String nodePath) throws KeeperException, InterruptedException {
        Stat stat = zookeeper.exists(nodePath, true);
        if (stat == null) {
            System.out.println(serverName + " 成为领导者！");
            isLeader = true;
        }
    }

    @Override
    public void process(WatchedEvent event) {
        if (event.getType() == Event.EventType.NodeDeleted) {
            try {
                watchNode(event.getPath());
            } catch (KeeperException | InterruptedException e) {
                e.printStackTrace();
            }
        }
        if (event.getType() == Event.EventType.NodeDataChanged) {
            try {
                System.out.println(serverName + " receive " + new String(zookeeper.getData(event.getPath(), true, null)));
            } catch (KeeperException e) {
                throw new RuntimeException(e);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
    }

    public static void main(String[] args) throws Exception {
        String connectString = "192.168.20.34:2181";
        String serverName = "Server-" + args[0];
//        String serverName = "Server-1";
        CountDownLatch countDownLatch = new CountDownLatch(1);
        LeaderElection election = new LeaderElection(connectString, serverName, countDownLatch);
        election.waitConnection(countDownLatch);
        election.startElection();
        Thread.sleep(Long.MAX_VALUE);
    }
}
